package dd.lo.ml;

import dd.lo.HelloCV;
import dd.lo.ml.enums.ImgLabel;
import org.apache.commons.collections4.CollectionUtils;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.core.MatOfFloat;
import org.opencv.core.TermCriteria;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;
import org.opencv.ml.Ml;
import org.opencv.ml.SVM;
import org.opencv.ml.TrainData;

import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ImgClassificationTrainer {

    private static final String ROOT = HelloCV.class.getResource("/").getPath();

    static {
        System.load(ROOT + "libopencv_java481.dylib");
    }

    public static void main(String[] args) {
        String rootPath = "/Users/kwok/Downloads/img_classification_train_data";
        Map<Integer, List<MatOfFloat>> trainingData = new HashMap<>(ImgLabel.values().length);
        int rowCnt = 0, colLen = 0;
        for (ImgLabel imgLabel : ImgLabel.values()) {
            File dir = new File(rootPath + "/" + imgLabel.getName());
            if (!dir.exists() || null == dir.listFiles()) {
                System.out.printf("文件夹【%s】不存在\n", dir.getAbsolutePath());
                continue;
            }
            System.out.printf("找到训练数据《%s》，开始遍历下面的文件...\n", imgLabel.getName());
            List<MatOfFloat> imgFeatures = new ArrayList<>(dir.listFiles().length);
            trainingData.put(imgLabel.getLabel(), imgFeatures);
            for (File file : dir.listFiles()) {
                if (null == file) continue;
                if (!file.getName().endsWith(".jpeg")) continue;
                System.out.printf("读取图片《%s》并计算图片特征...\n", file.getName());
                Mat grayImg = Imgcodecs.imread(file.getAbsolutePath(), Imgcodecs.IMREAD_GRAYSCALE);
                MatOfFloat imgHogFeature = FeatureExtraction.computeGrayImgHogFeature(grayImg);
                colLen = imgHogFeature.height();
                imgFeatures.add(imgHogFeature);
                rowCnt++;
            }
        }
        Mat trainMat = new Mat(rowCnt, colLen, CvType.CV_32FC1);
        Mat labelMat = new Mat(rowCnt, 1, CvType.CV_32SC1);
        int rowId = 0;
        for (Map.Entry<Integer, List<MatOfFloat>> entry : trainingData.entrySet()) {
            for (MatOfFloat trainDatum : entry.getValue()) {
                labelMat.put(rowId, 0, entry.getKey());
                trainMat.put(rowId, 0, trainDatum.toArray());
                rowId++;
            }
        }
        TrainData trainData = TrainData.create(trainMat, Ml.ROW_SAMPLE, labelMat);
        // 创建SVM对象并设置参数
        System.out.println("训练数据已经准备好，开始训练模型...");
        SVM svm = SVM.create();
        svm.setType(SVM.C_SVC); // 设置SVM类型为C_SVC（分类）
        svm.setKernel(SVM.LINEAR); // 设置核函数为线性核函数
        svm.setTermCriteria(new TermCriteria(TermCriteria.MAX_ITER, 100, 1e-6));
        svm.train(trainData);
        //保存SVM模型
        System.out.println("训练完成，保存模型...");
        svm.save(ROOT + "svm_classification_v_1_0.xml");
    }
}
